//
// Created by ttai on 2021/12/20.
//
#include "sm4.h"

string padding_128bit(string m) {
    int len = m.size();
    int r = len % 128;
    if (r != 0)m.insert(m.size(), 128 - r, '0');
    return m;
}

string Tao(string A) {
    string a1 = A.substr(0, 8);
    string a2 = A.substr(8, 8);
    string a3 = A.substr(16, 8);
    string a4 = A.substr(24, 8);
    string b = "";
    int x = stoi(a1.substr(0, 4), nullptr, 2);
    int y = stoi(a1.substr(4, 4), nullptr, 2);
    b += hex_bin(SboxTable[x][y]);
    x = stoi(a2.substr(0, 4), nullptr, 2);
    y = stoi(a2.substr(4, 4), nullptr, 2);
    b += hex_bin(SboxTable[x][y]);
    x = stoi(a3.substr(0, 4), nullptr, 2);
    y = stoi(a3.substr(4, 4), nullptr, 2);
    b += hex_bin(SboxTable[x][y]);
    x = stoi(a4.substr(0, 4), nullptr, 2);
    y = stoi(a4.substr(4, 4), nullptr, 2);
    b += hex_bin(SboxTable[x][y]);

    return b;
}

string L(string b) {
    bitset<32> B(b);
    bitset<32> C = B ^ CROL32(B, 2) ^ CROL32(B, 10) ^ CROL32(B, 18) ^ CROL32(B, 24);
    return C.to_string();
}

string L_(string b) {
    bitset<32> B(b);
    bitset<32> ret = B ^ CROL32(B, 13) ^ CROL32(B, 23);
    return ret.to_string();
}

string T(string s) {
    string ret;
    ret = Tao(s);
    return L(ret);
}

string T_(string s) {
    string ret;
    ret = Tao(s);
    return L_(ret);
}


string F(string x0, string x1, string x2, string x3, string rk) {
    bitset<32> X0(x0);
    bitset<32> X1(x1);
    bitset<32> X2(x2);
    bitset<32> X3(x3);
    bitset<32> RK(rk);
    bitset<32> tmp = X1 ^ X2 ^ X3 ^ RK;
    string ret = T(tmp.to_string());
    bitset<32> result = X0 ^ bitset<32>(ret);
    return result.to_string();
}

int key_expand(string MK, vector<string> &rk_list) {
    string mk0 = MK.substr(0, 32);
    string mk1 = MK.substr(32, 32);
    string mk2 = MK.substr(64, 32);
    string mk3 = MK.substr(96, 32);
    bitset<32> MK0(mk0);
    bitset<32> MK1(mk1);
    bitset<32> MK2(mk2);
    bitset<32> MK3(mk3);
    vector<string> k_list;
    k_list.push_back((MK0 ^ bitset<32>(hex_bin(FK0))).to_string());
    k_list.push_back((MK1 ^ bitset<32>(hex_bin(FK1))).to_string());
    k_list.push_back((MK2 ^ bitset<32>(hex_bin(FK2))).to_string());
    k_list.push_back((MK3 ^ bitset<32>(hex_bin(FK3))).to_string());
    for (int i = 0; i < 32; i++) {
        bitset<32> tmp = (bitset<32>(k_list[i + 1]) ^ bitset<32>(k_list[i + 2]) ^ bitset<32>(k_list[i + 3]) ^
                          bitset<32>(hex_bin(CK[i])));
        bitset<32> rk = bitset<32>(k_list[i]) ^ bitset<32>(T_(tmp.to_string()));
        k_list.push_back(rk.to_string());
        rk_list.push_back(rk.to_string());
    }
    return 1;
}

string encrypt(string x, vector<string> rk_list) {
    string x0 = x.substr(0, 32);
    string x1 = x.substr(32, 32);
    string x2 = x.substr(64, 32);
    string x3 = x.substr(96, 32);
    vector<string> X_list;
    X_list.push_back(x0);
    X_list.push_back(x1);
    X_list.push_back(x2);
    X_list.push_back(x3);
    for (int i = 0; i < 32; i++) {
        string Xtmp = F(X_list[i], X_list[i + 1], X_list[i + 2], X_list[i + 3], rk_list[i]);
        X_list.push_back(Xtmp);
    }
    string ret = X_list[35] + X_list[34] + X_list[33] + X_list[32];
    return bin_hex(ret);
}

string decrypt(string x, vector<string> rk_list) {
    reverse(rk_list.begin(), rk_list.end());  //轮密钥反转
    string x0 = x.substr(0, 32);
    string x1 = x.substr(32, 32);
    string x2 = x.substr(64, 32);
    string x3 = x.substr(96, 32);
    vector<string> X_list;
    X_list.push_back(x0);
    X_list.push_back(x1);
    X_list.push_back(x2);
    X_list.push_back(x3);
    for (int i = 0; i < 32; i++) {
        string Xtmp = F(X_list[i], X_list[i + 1], X_list[i + 2], X_list[i + 3], rk_list[i]);
        X_list.push_back(Xtmp);
    }
    string ret = X_list[35] + X_list[34] + X_list[33] + X_list[32];
    return bin_hex(ret);
}

string sm4_encrypt(string s, string mk, int mode, string iv) {
    string s_bin = string_bin(s);
    string s_bin_128bit = padding_128bit(s_bin);
    string mk_bin = string_bin(mk);
    vector<string> rk_list;
    key_expand(mk_bin, rk_list);
    string iv_bin = string_bin(iv);
    switch (mode) {
        case 1:
            return _ecb_encrypt(s_bin_128bit, rk_list);
        case 2:
            return _cbc_encrypt(s_bin_128bit, rk_list, iv_bin);
        case 3:
            return _cfb_encrypt(s_bin_128bit, rk_list, iv_bin);
        case 4:
            return _ofb_encrypt(s_bin_128bit, rk_list, iv_bin);
    }

}

string sm4_decrypt(string s, string mk, int mode, string iv) {

    string s_bin = hex_bin(s);
    string mk_bin = string_bin(mk);
    vector<string> rk_list;
    key_expand(mk_bin, rk_list);
    string iv_bin = string_bin(iv);
    switch (mode) {
        case 1:
            return _ecb_decrypt(s_bin, rk_list);
        case 2:
            return _cbc_decrypt(s_bin, rk_list, iv_bin);
        case 3:
            return _cfb_decrypt(s_bin, rk_list, iv_bin);
        case 4:
            return _ofb_decrypt(s_bin, rk_list, iv_bin);
    }

}

string _ecb_encrypt(string s, vector<string> rk_list) {
    string cipher = "";
    for (int i = 0; i < s.size(); i += 128) {
        string tmp = s.substr(i, 128);
        cipher += encrypt(tmp, rk_list);
    }
    return cipher;
}

string _ecb_decrypt(string s, vector<string> rk_list) {
    string plain = "";
    for (int i = 0; i < s.size(); i += 128) {
        string tmp = s.substr(i, 128);
        plain += decrypt(tmp, rk_list);
    }
    string ret = hex_string(plain);
    return ret;
}

string _cbc_encrypt(string s, vector<string> rk_list, string iv) {

    bitset<128> tmp_iv(iv);
    string cipher = "";
    for (int i = 0; i < s.size(); i += 128) {
        bitset<128> tmp(s.substr(i, 128));
        tmp ^= tmp_iv;
        string tmp_cipher = encrypt(tmp.to_string(), rk_list);
        tmp_iv = bitset<128>(hex_bin(tmp_cipher));
        cipher += tmp_cipher;
    }
    return cipher;
}

string _cbc_decrypt(string s, vector<string> rk_list, string iv) {
    bitset<128> tmp_iv(iv);
    string plain = "";
    for (int i = 0; i < s.size(); i += 128) {
        string str = s.substr(i, 128);
        string str_plain = decrypt(str, rk_list);
        string tmp_plain = (bitset<128>(hex_bin(str_plain)) ^ tmp_iv).to_string();
        plain += bin_hex(tmp_plain);
        tmp_iv = bitset<128>(str);
    }
    return hex_string(plain);
}

string _cfb_encrypt(string s, vector<string> rk_list, string iv) {
    string tmp_iv = iv;
    string cipher = "";
    for (int i = 0; i < s.size(); i += 128) {
        string str = s.substr(i, 128);
        string iv_cipher = encrypt(tmp_iv, rk_list);
        string tmp_cipher = (bitset<128>(str) ^ bitset<128>(hex_bin(iv_cipher))).to_string();
        cipher += bin_hex(tmp_cipher);
        tmp_iv = tmp_cipher;
    }
    return cipher;
}

string _cfb_decrypt(string s, vector<string> rk_list, string iv) {
    string tmp_iv = iv;
    string plain = "";
    for (int i = 0; i < s.size(); i += 128) {
        string str = s.substr(i, 128);
        string iv_cipher = encrypt(tmp_iv, rk_list);
        string tmp_plain = (bitset<128>(str) ^ bitset<128>(hex_bin(iv_cipher))).to_string();
        plain += bin_hex(tmp_plain);
        tmp_iv = str;
    }
    return hex_string(plain);
}

string _ofb_encrypt(string s, vector<string> rk_list, string iv) {
    string tmp_iv = iv;
    string cipher = "";
    for (int i = 0; i < s.size(); i += 128) {
        string str = s.substr(i, 128);
        tmp_iv = encrypt(tmp_iv, rk_list);
        string tmp_cipher = (bitset<128>(str) ^ bitset<128>(hex_bin(tmp_iv))).to_string();
        cipher += bin_hex(tmp_cipher);
        tmp_iv = hex_bin(tmp_iv);
    }
    return cipher;
}

string _ofb_decrypt(string s, vector<string> rk_list, string iv) {
    string tmp_iv = iv;
    string plain = "";
    for (int i = 0; i < s.size(); i += 128) {
        string str = s.substr(i, 128);
        tmp_iv = encrypt(tmp_iv, rk_list);
        string tmp_plain = (bitset<128>(str) ^ bitset<128>(hex_bin(tmp_iv))).to_string();
        plain += bin_hex(tmp_plain);
        tmp_iv = hex_bin(tmp_iv);
    }
    return hex_string(plain);
}
